import torch

from .plugins_base import PluginBase
from trainer.trainer import PluginType, TrainContext
from utils.segmentation import SegmentationMetric


class TrainingMetricsPlugin(PluginBase):
    plugin_hooks = {
        PluginType.EPOCH_END: "log_metrics",
        PluginType.BATCH_END: "update_miou",
    }

    def __init__(self, num_classes=19, ignore_index=255):
        self.train_loss = []
        self.train_miou = []
        self.train_mpa = []
        self.miou_cal = SegmentationMetric(num_classes=num_classes, ignore_index=ignore_index)
    
    def update_miou(self, ctx: TrainContext):
        predictions = torch.argmax(ctx.outputs['out'], dim=1)
        self.miou_cal.update(predictions, ctx.labels)
        print(f'\rBatch: {ctx.batch+1}/{len(ctx.train_loader)}', end='', flush=True)
        if ctx.batch + 1 == len(ctx.train_loader):
            print('\r', end='', flush=True)

    def log_metrics(self, ctx: TrainContext):
        avg_loss = ctx.avg_loss
        time_spent = ctx.workspace.get("train_time", 0.0)

        metrics = self.miou_cal.evaluate()
        avg_miou = metrics["mIoU"]
        avg_mpa = metrics["mPA"]

        self.train_loss.append(avg_loss)
        self.train_miou.append(avg_miou)
        self.train_mpa.append(avg_mpa)

        msg = {
            "mode": "train",
            "epoch": ctx.epoch + 1,
            "lr": round(ctx.optimizer.param_groups[0]['lr'], 6),
            "time": round(time_spent, 6),
            "loss": round(avg_loss, 6),
            "acc": round(avg_mpa, 6),
            "miou": round(avg_miou, 6),
        }

        if self.check_key(ctx.workspace, "logger"):
            ctx.workspace["logger"](str(msg))
